package org.budo.mongo.jdbc.driver.executor;

import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

import org.bson.Document;
import org.budo.jdbc.driver.JdbcUrl;
import org.budo.mongo.jdbc.driver.BudoMongoJdbcPreparedStatement;
import org.budo.mongo.jdbc.driver.statement.MongoStatement;
import org.budo.mongo.jdbc.driver.statement.StatementDocument;
import org.budo.mongo.jdbc.resultset.impl.MongoMapArrayResultSet;
import org.budo.support.lang.util.ListUtil;
import org.budo.support.lang.util.MapUtil;

import com.mongodb.DB;
import com.mongodb.MongoClient;
import com.mongodb.MongoClientOptions;
import com.mongodb.MongoCommandException;
import com.mongodb.ServerAddress;
import com.mongodb.client.FindIterable;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.MongoDatabase;
import com.mongodb.client.MongoIterable;
import com.mongodb.client.result.DeleteResult;
import com.mongodb.client.result.UpdateResult;

import lombok.Getter;
import lombok.NoArgsConstructor;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;

/**
 * @author lmw
 */
@Slf4j
@Getter
@Setter
@NoArgsConstructor
public class DefaultMongoStatementExecutor implements MongoStatementExecutor {
    private static final int SOCKET_TIMEOUT = 15000;

    private static final int CONNECT_TIMEOUT = 5000;

    private static final String EVAL_FUNCTION_SCRIPT = "var cursorToArray = function (cursor) {" //
            + "    var result = [];" //
            + "    while (cursor.hasNext()) {" //
            + "        result.push(cursor.next());" //
            + "    }" //
            + "    return result;" //
            + "};" //
            + "" //
            + "var copyToTable = function (cursor, tableTo) {" //
            + "    while (cursor.hasNext()) {" //
            + "        db[tableTo].insert(cursor.next());" //
            + "    }" //
            + "};";

    private JdbcUrl jdbcUrl;

    private MongoClient _mongoClient;

    private MongoDatabase _mongoDatabase;

    /**
     * 延迟初始化
     */
    private MongoClient mongoClient() {
        if (null != this._mongoClient) {
            return this._mongoClient;
        }

        synchronized (this) {
            if (null != this._mongoClient) {
                return this._mongoClient;
            }

            ServerAddress serverAddress = new ServerAddress(this.getJdbcUrl().getHost(), this.getJdbcUrl().getPort());
            MongoClientOptions mongoClientOptions = MongoClientOptions.builder() //
                    .connectTimeout(CONNECT_TIMEOUT) // 超时
                    .socketTimeout(SOCKET_TIMEOUT) // 超时
                    .build();

            return this._mongoClient = new MongoClient(serverAddress, mongoClientOptions);
        }
    }

    /**
     * 延迟初始化
     */
    private MongoDatabase mongoDatabase() {
        if (null != this._mongoDatabase) {
            return this._mongoDatabase;
        }

        synchronized (this) {
            if (null != this._mongoDatabase) {
                return this._mongoDatabase;
            }

            return this._mongoDatabase = this.mongoClient().getDatabase(this.getJdbcUrl().getCatalog());
        }
    }

    public DefaultMongoStatementExecutor(JdbcUrl jdbcUrl) {
        this.setJdbcUrl(jdbcUrl);
    }

    private List<Document> mapToDocument(List<Map<String, Object>> insertDocuments) {
        List<Document> documents = new ArrayList<Document>();
        for (Map<String, Object> insert : insertDocuments) {
            documents.add(new Document(insert));
        }
        return documents;
    }

    @Override
    public List<?> eval(String evalScript) {
        try {
            String dbName = this.mongoDatabase().getName();

            @SuppressWarnings("deprecation")
            DB db = this.mongoClient().getDB(dbName);

            String code = EVAL_FUNCTION_SCRIPT + evalScript;
            return (List<?>) db.eval(code);
        } catch (MongoCommandException e) {
            throw new RuntimeException("#66 evalScript=" + evalScript, e);
        }
    }

    @Override
    public Iterator<?> aggregate(String collectionName, List<Map<String, Object>> aggregatePipeline) {
        List<Document> pipeline = this.mapToDocument(aggregatePipeline);

        MongoCollection<Document> collection = this.mongoDatabase().getCollection(collectionName);
        return collection.aggregate(pipeline).iterator();
    }

    @Override
    public long count(String collectionName, Map<String, Object> where) {
        MongoCollection<Document> collection = this.mongoDatabase().getCollection(collectionName);
        if (null == where || where.isEmpty()) {
            return collection.count();
        }

        return collection.count(new Document(where));
    }

    @Override
    public Integer delete(String collectionName, Map<String, Object> whereDocument) {
        MongoCollection<Document> collection = this.mongoDatabase().getCollection(collectionName);
        DeleteResult deleteResult = collection.deleteMany(new Document(whereDocument));
        return (int) deleteResult.getDeletedCount();
    }

    @Override
    public void insert(String collectionName, List<Map<String, Object>> insertDocuments, PreparedStatement preparedStatement) {
        List<Document> documents = this.mapToDocument(insertDocuments);

        // 插入数据
        MongoCollection<Document> collection = this.mongoDatabase().getCollection(collectionName);
        collection.insertMany(documents);

        if (log.isDebugEnabled()) {
            log.debug("#126 after insertMany, collectionName=" + collectionName + ", size=" + ListUtil.size(documents));
        }

        this.handleAutoGeneratedKeys(preparedStatement, documents);
    }

    private void handleAutoGeneratedKeys(PreparedStatement preparedStatement, List<Document> documents) {
        BudoMongoJdbcPreparedStatement mongoPreparedStatement = (BudoMongoJdbcPreparedStatement) preparedStatement;

        if (Statement.RETURN_GENERATED_KEYS != mongoPreparedStatement.getAutoGeneratedKeys()) {
            return; // 不需要返回主键
        }

        // 收集 setGeneratedKeys
        int size = documents.size();
        Map<?, ?>[] valueArray = new Map<?, ?>[size];
        for (int i = 0; i < size; i++) {
            Document document = documents.get(i);
            Object _id = document.get(MongoStatement._ID);
            valueArray[i] = MapUtil.stringObjectMap(MongoStatement._ID, _id);
        }

        ResultSet resultSet = new MongoMapArrayResultSet(new String[] { MongoStatement._ID }, valueArray);

        mongoPreparedStatement.setGeneratedKeys(resultSet);
    }

    @Override
    public Integer update(String collectionName, Map<String, Object> where, Map<String, Object> update) {
        MongoCollection<Document> collection = this.mongoDatabase().getCollection(collectionName);
        UpdateResult updateResult = collection.updateMany(new Document(where), new Document(update));
        return (int) updateResult.getModifiedCount();
    }

    @Override
    public Iterator<?> find(String collectionName, StatementDocument statementDocument) {
        MongoCollection<Document> collection = this.mongoDatabase().getCollection(collectionName);

        Map<String, Object> where = statementDocument.getWhere();

        FindIterable<Document> findIterable;
        if (null == where) {
            findIterable = collection.find();
        } else {
            findIterable = collection.find(new Document(where));
        }

        Map<String, Object> sort = statementDocument.getSort();
        if (null != sort && !sort.isEmpty()) {
            findIterable = findIterable.sort(new Document(sort));
        }

        if (null != statementDocument.getSkip()) {
            findIterable = findIterable.skip(statementDocument.getSkip());
        }

        if (null != statementDocument.getLimit()) {
            findIterable = findIterable.limit(statementDocument.getLimit());
        }

        return findIterable.iterator();
    }

    @Override
    public Iterator<?> listDatabaseNames() {
        MongoIterable<String> listDatabaseNames = this.mongoClient().listDatabaseNames();
        return listDatabaseNames.iterator();
    }

    @Override
    public Iterator<?> listCollectionNames() {
        MongoIterable<String> listCollectionNames = this.mongoDatabase().listCollectionNames();
        return listCollectionNames.iterator();
    }

    @Override
    public void close() {
        if (null != this._mongoClient) {
            this._mongoClient.close();
        }
    }
}